import torch
import numpy as np


def compute_optimal_transport(args, M, epsilon=1e-6):
    #r is the P we discussed in paper r.shape = n_runs x total_queries, all entries = 1
    r = torch.ones(1, M.shape[0])
    #r = r * weights
    #c = torch.ones(1, M.shape[1]) * int(M.shape[0]/M.shape[1])
    if args.unlabel == 0:
        unlabel = 15
    else:
        unlabel = args.unlabel
    unlabel = np.array(np.repeat(float(unlabel), args.way))
    c = torch.FloatTensor(unlabel)
    idx = np.where(c.detach().cpu().numpy() <= 0)

    M = M.cuda()
    r = r.cuda()
    c = c.cuda()
    M = torch.unsqueeze(M, dim=0)
    n_runs, n, m = M.shape
    P = M

    u = torch.zeros(n_runs, n).cuda()
    maxiters = 1000
    iters = 1
    # for i in range(opt.sinkhorn_iter):
    for i in range(1):
        # P = torch.pow(P, opt.T)
        P = torch.pow(P, 3)
        while torch.max(torch.abs(u - P.sum(2))) > epsilon:
            u = P.sum(2)
            P *= (r / u).view((n_runs, -1, 1))
            P *= (c / P.sum(1)).view((n_runs, 1, -1))
            if len(idx[0]) > 0:
                P[P != P] = 0
            if iters == maxiters:
                break
            iters = iters + 1
    P = torch.squeeze(P).detach().cpu().numpy()
    best_per_class = np.argmax(P, 0)
    if M.shape[1]==1:
        P = np.expand_dims(P, axis=0)
    labels = np.argmax(P, 1)
    return P, labels, best_per_class